import glob
import math
import os
import xml.etree.ElementTree as ET

import cv2
import numpy as np


def order_points_new(pts):
    # sort the points based on their x-coordinates
    xSorted = pts[np.argsort(pts[:, 0]), :]

    # grab the left-most and right-most points from the sorted
    # x-roodinate points
    leftMost = xSorted[:2, :]
    rightMost = xSorted[2:, :]
    if leftMost[0, 1] != leftMost[1, 1]:
        leftMost = leftMost[np.argsort(leftMost[:, 1]), :]
    else:
        leftMost = leftMost[np.argsort(leftMost[:, 0])[::-1], :]
    (tl, bl) = leftMost
    if rightMost[0, 1] != rightMost[1, 1]:
        rightMost = rightMost[np.argsort(rightMost[:, 1]), :]
    else:
        rightMost = rightMost[np.argsort(rightMost[:, 0])[::-1], :]
    (tr, br) = rightMost
    # return the coordinates in top-left, top-right,
    # bottom-right, and bottom-left order
    return np.array([tl, tr, br, bl], dtype="float32")


def polygenInBox(pts, box):
    x1, y1, x2, y2 = box
    for pt in pts:
        if x1 <= pt[0] <= x2 and \
                (y1 <= pt[1] <= y2):
            continue
        else:
            return False
    return True


def get_min_rect_4points(roi):
    center = roi[0]
    (w, h) = roi[1]
    min_rect_angle = roi[2]
    if h > w:
        min_rect_angle = min_rect_angle + 180

    angelPi = (min_rect_angle / 180) * math.pi
    x = center[0]
    y = center[1]
    x1 = x + (w / 2) * math.cos(angelPi) - (h / 2) * math.sin(angelPi)
    y1 = y + (w / 2) * math.sin(angelPi) + (h / 2) * math.cos(angelPi)

    x2 = x + (w / 2) * math.cos(angelPi) + (h / 2) * math.sin(angelPi)
    y2 = y + (w / 2) * math.sin(angelPi) - (h / 2) * math.cos(angelPi)

    x3 = x - (w / 2) * math.cos(angelPi) + (h / 2) * math.sin(angelPi)
    y3 = y - (w / 2) * math.sin(angelPi) - (h / 2) * math.cos(angelPi)

    x4 = x - (w / 2) * math.cos(angelPi) - (h / 2) * math.sin(angelPi)
    y4 = y - (w / 2) * math.sin(angelPi) + (h / 2) * math.cos(angelPi)
    return np.array([
        [x1, y1],
        [x2, y2],
        [x3, y3],
        [x4, y4],
    ])


def find_polygen_in_box(box, pts_str_color_list):
    for pts, str, color in pts_str_color_list:
        if polygenInBox(pts, box):
            return pts, str, color
    return None, None, None


class ScopeCarStru:
    def __init__(self, file_name):
        self.car_coor = None
        self.lp_pts = None
        self.lp_color = None
        self.relative_lp_pts = np.zeros([4, 2])
        self.ocr = None
        self.file_name = file_name

        self.lp_pts_prob = None
        self.car_prob = None
        self.cls_name = None
        self.cls_ID = None
        self.warpped_im = None

    def make_relateve_lp_pts(self):
        x1, y1, x2, y2 = self.car_coor
        self.car_w = float(x2 - x1)
        self.car_h = float(y2 - y1)
        xc = (x1 + x2) / 2
        yc = (y1 + y2) / 2
        self.relative_lp_pts = (self.lp_pts - np.array([x1, y1]).reshape(1, 2)) / np.array([self.car_w, self.car_h]).reshape(1, 2)
        # 相对于中心点


class FulPicStru:
    def __init__(self):
        self.cars = []
        self.file_name = None
        self.width = None
        self.height = None


class MetaLpContainer:
    def make_all_pic_strus(self):
        for i in range(len(self.all_files)):
            pic_stru = FulPicStru()
            pic_stru.file_name = self.all_files[i]
            pic_stru.width = self.all_whs[i][0]
            pic_stru.height = self.all_whs[i][1]

            for car_coor in self.all_cars_pts[i]:
                a_car = ScopeCarStru(pic_stru.file_name)
                a_car.car_coor = car_coor
                lp_pts, ocr, color = find_polygen_in_box(car_coor, self.all_lp_pts_str_color[i])
                if lp_pts is not None:
                    a_car.lp_pts = lp_pts
                    a_car.ocr = ocr
                    a_car.lp_color = color

                    a_car.make_relateve_lp_pts()
                pic_stru.cars.append(a_car)
                self.all_expanded_car_strus.append(a_car)
            self.all_pic_strus.append(pic_stru)

    def __init__(self, root="/home/data/782"):
        self.root = root
        self.all_files = []  # pics
        self.all_whs = []
        self.all_cars_pts = []  # [[pic 1 car1 ,pic 1 car2],[pic 2 car1 ,pic 2 car2]]
        self.all_lp_pts_str_color = []  # [[pic 1 lp1 ,pic 1 lp2],[]]

        self.all_pic_strus = []
        self.all_expanded_car_strus = []

        self.anno_files = glob.glob(root + "/*.xml")

        for anno in self.anno_files:
            tree = ET.parse(anno)
            root = tree.getroot()

            filename = root.find("filename").text

            width = int(root.find("size").find("width").text)
            height = int(root.find("size").find("height").text)

            pic_car_pts = []
            pic_lp_pts_str_color = []

            for child in root.iter('object'):
                name = child.find("name").text
                x1 = float(child.find("bndbox").find("xmin").text)
                y1 = float(child.find("bndbox").find("ymin").text)
                x2 = float(child.find("bndbox").find("xmax").text)
                y2 = float(child.find("bndbox").find("ymax").text)
                print("{} box:{} {} {} {} {}".format(filename, name, x1, y1, x2, y2))
                pic_car_pts.append([x1, y1, x2, y2])

            for polygen in root.iter("polygon"):
                clsname = polygen.find("class").text
                assert clsname == "plate"
                points = polygen.find("points").text
                print("points text:{}".format(points))
                points = [p for p in points.split(";")]
                print("points text split ; :{}".format(points))
                points = [[float(p2.split(",")[0]), float(p2.split(",")[1])] for p2 in points]
                print("points text split , :{}".format(points))
                if len(points) > 4:
                    points = np.array(points).reshape(len(points), 2)
                    rect = cv2.minAreaRect(points.astype(np.int32).reshape(-1, 1, 2))
                    points = get_min_rect_4points(rect)
                else:
                    points = np.array(points).reshape(4, 2)

                points = order_points_new(points)
                x1 = points[0, 0]
                y1 = points[0, 1]
                x2 = points[1, 0]
                y2 = points[1, 1]
                x3 = points[2, 0]
                y3 = points[2, 1]
                x4 = points[3, 0]
                y4 = points[3, 1]

                assert (x2 > x1)
                assert (y3 > y2)
                assert (x3 > x4)
                assert (y1 < y4)

                attrs = [it for it in polygen.iter("attribute")]
                assert len(attrs) == 2
                lp_str = attrs[0].find("value").text
                color = attrs[1].find("value").text
                print("{} lp:{} color:{} ps:{}".format(filename, lp_str, color, points))
                pic_lp_pts_str_color.append([points, lp_str, color])

            self.all_files.append(filename)
            self.all_whs.append([width, height])
            self.all_cars_pts.append(pic_car_pts)
            self.all_lp_pts_str_color.append(pic_lp_pts_str_color)

        # TODO NOTE: lp str has None label
        self.make_all_pic_strus()
        print("LOAD DONE")
        return

    def get_car_im_lp_pts_ocr_color(self, idx):
        a_car = self.all_expanded_car_strus[idx]
        x1, y1, x2, y2 = a_car.car_coor
        bgr_im = cv2.imread(os.path.join(self.root, a_car.file_name))
        # car_bgr_im = bgr_im[int(y1):int(y2), int(x1):int(x2), :]
        # return car_bgr_im, a_car.relative_lp_pts, a_car.ocr, a_car.lp_color

        # full_bgr , lp_pts = (xp - x_car) / car_w
        return bgr_im, (x1, y1, x2, y2), a_car.relative_lp_pts, a_car.ocr, a_car.lp_color

    def get_a_pic_cars(self, idx):
        pic_stru = self.all_pic_strus[idx]
        bgr_im = cv2.imread(os.path.join(self.root, pic_stru.file_name))
        car_boxes = np.zeros([len(pic_stru.cars), 4])
        car_ocrs = []
        car_lp_colors = []
        car_lp_ptses = np.zeros([len(pic_stru.cars), 4, 2])

        for i, a_car in enumerate(pic_stru.cars):
            x1, y1, x2, y2 = a_car.car_coor
            car_boxes[i] = np.array([x1, y1, x2, y2]).reshape(1, 4)
            car_ocrs.append(a_car.ocr)
            car_lp_colors.append(a_car.lp_color)
            car_lp_ptses[i] = a_car.relative_lp_pts
        return bgr_im, car_boxes, car_lp_ptses, car_ocrs, car_lp_colors


if __name__ == '__main__':
    con = MetaLpContainer()
